/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: pingmesh 对外接口实现
 */

#include "network/hccp_ping.h"
#include "ping_mesh.h"
#include "hccl_ip_address.h"
#include "dlra_function.h"
#include "adapter_rts_common.h"
#include "hccn_rping.h"

#ifdef __cplusplus
extern "C" {
#endif // __cplusplus

using namespace hccl;
constexpr u32 BUFFER_SIZE_UNIT = 4096;
constexpr u32 NPU_NUM_MAX = 32768;
constexpr u32 NPU_NUM_MIN = 128;
constexpr u32 NPU_NUM_MAX_BITWIDTH = 128;
constexpr u32 LINK_TYPE_MODE_ROCE = 3;
constexpr u32 RPING_RESULT_STATE_VALID = 2;
constexpr u32 TARGET_NUM_MAX = 16;

HccnResult HccnRpingInit(uint32_t devLogicId, HccnRpingInitAttr *initAttr, HccnRpingCtx *rpingCtx)
{
    // 校验入参
    CHK_PRT_RET(initAttr == nullptr, HCCL_ERROR("[HccnRpingInit]initAttr is null."), HCCN_E_PARA);
    CHK_PRT_RET(initAttr->mode >= HCCN_RPING_MODE_RESERVED,
        HCCL_ERROR("[HccnRpingInit]LinkMode[%d] not support.", initAttr->mode), HCCN_E_PARA);
    CHK_PRT_RET(rpingCtx == nullptr, HCCL_ERROR("[HccnRpingInit]rpingCtx is null."), HCCN_E_PARA);
    // 初始化ra接口
    CHK_PRT_RET(DlRaFunction::GetInstance().DlRaFunctionInit() != HCCL_SUCCESS,
        HCCL_ERROR("[HccnRpingInit]dlratunction failed."), HCCN_E_FAIL);
    // 获取device id
    s32 currDevLogicId = 0;
    HcclResult ret = hrtGetDeviceRefresh(&currDevLogicId);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingInit]cannot get device logic id."), HCCN_E_FAIL);

    // 判断devLogicId与pingmesh所记录的是否一致
    if (devLogicId != static_cast<u32>(currDevLogicId)) {
        HCCL_ERROR("[HccnRpingInit]Input device logicId[%d] don't match current logicId[%d].",
            devLogicId, currDevLogicId);
        return HCCN_E_FAIL;
    }
    // 构造一个pingmesh实例, 用于存放初始化的pingmesh信息
    PingMesh *rping = new (std::nothrow) PingMesh();
    CHK_PRT_RET(rping == nullptr, HCCL_ERROR("[HccnRpingInit]rping alloc failed."), HCCN_E_MEM);
    // 利用用户输入的ip地址构造HcclIpAddress类
    HcclIpAddress ipAddr = HcclIpAddress(std::string(initAttr->ipAddr));
    // npuNum必须为2的整次幂
    u32 npuNum = (initAttr->npuNum > NPU_NUM_MIN) ? initAttr->npuNum : NPU_NUM_MIN;
    npuNum = (npuNum > NPU_NUM_MAX) ? (NPU_NUM_MAX - 1) : (npuNum - 1);
    for (u32 i = 0; i < NPU_NUM_MAX_BITWIDTH; i++) {
        npuNum |= (npuNum >> 1);
    }
    npuNum++;
    npuNum = npuNum < NPU_NUM_MIN ? NPU_NUM_MIN : npuNum;
    // buffer_size必须为4k的倍数
    u32 buffer_size = 0;
    if (initAttr->bufferSize > 0) {
        buffer_size = ((initAttr->bufferSize - 1) / BUFFER_SIZE_UNIT + 1) * BUFFER_SIZE_UNIT;
    }

    // 初始化pingmesh实例
    ret = rping->HccnRpingInit(devLogicId, LINK_TYPE_MODE_ROCE, ipAddr, initAttr->port, npuNum, buffer_size,
        initAttr->sl, initAttr->tc);
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[HccnRpingInit] Pingmesh init failed.");
        delete rping; // 初始化失败释放内存
        return HCCN_E_FAIL;
    }

    // 记录pingmesh指针
    *rpingCtx = rping;
    return HCCN_SUCCESS;
}
 
HccnResult HccnRpingDeinit(HccnRpingCtx rpingCtx)
{
    // 初始化ra接口
    CHK_PRT_RET(DlRaFunction::GetInstance().DlRaFunctionInit() != HCCL_SUCCESS,
        HCCL_ERROR("[HccnRpingDeinit]dlratunction failed."), HCCN_E_FAIL);
    // 校验指针，并将其转换为pingmesh指针
    CHK_PRT_RET(rpingCtx == nullptr, HCCL_ERROR("[HccnRpingDeinit]rpingCtx is null."), HCCN_E_PARA);
    PingMesh *rping = static_cast<PingMesh*>(rpingCtx);
    // 获取device id
    s32 devLogicId = 0;
    HcclResult ret = hrtGetDeviceRefresh(&devLogicId);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingDeinit]cannot get device logic id."), HCCN_E_FAIL);

    // 判断devLogicId与pingmesh所记录的是否一致
    if (devLogicId != rping->GetDeviceLogicId()) {
        HCCL_ERROR("[HccnRpingDeinit]curr devId[%d] don't match record logicId[%d].",
            devLogicId, rping->GetDeviceLogicId());
        return HCCN_E_PARA;
    }
 
    ret = rping->HccnRpingDeinit(static_cast<u32>(devLogicId));
    delete rping; // 先释放资源
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[HccnRpingDeinit]Device[%d] deinit fail", devLogicId);
        return HCCN_E_FAIL;
    }

    return HCCN_SUCCESS;
}
 
HccnResult HccnRpingAddTarget(HccnRpingCtx rpingCtx, uint32_t targetNum, HccnRpingTargetInfo *target)
{
    // 校验入参
    CHK_PRT_RET(target == nullptr, HCCL_ERROR("[HccnRpingAddTarget]target is null."), HCCN_E_PARA);
    CHK_PRT_RET(targetNum > TARGET_NUM_MAX,
        HCCL_ERROR("[HccnRpingAddTarget]targetNum[%u] is more than %u!", targetNum, TARGET_NUM_MAX), HCCN_E_PARA);
    // 校验指针，并将其转换为pingmesh指针
    CHK_PRT_RET(rpingCtx == nullptr, HCCL_ERROR("[HccnRpingAddTarget]rpingCtx is null."), HCCN_E_PARA);
    PingMesh *rping = static_cast<PingMesh*>(rpingCtx);
    // 获取device id
    s32 devLogicId = 0;
    HcclResult ret = hrtGetDeviceRefresh(&devLogicId);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingAddTarget]cannot get device logic id."), HCCN_E_FAIL);

    // 判断devLogicId与pingmesh所记录的是否一致
    if (devLogicId != rping->GetDeviceLogicId()) {
        HCCL_ERROR("[HccnRpingAddTarget]curr devId[%d] don't match record logicId[%d].",
            devLogicId, rping->GetDeviceLogicId());
        return HCCN_E_PARA;
    }
 
    // 将入参转换为内部接口可以使用的类型
    RpingInput *input = new (std::nothrow) RpingInput[targetNum];
    CHK_PRT_RET(input == nullptr, HCCL_ERROR("[HccnRpingAddTarget]memory alloc failed."), HCCN_E_MEM);
    for (uint32_t i = 0; i < targetNum; i++) {
        input[i].sip = HcclIpAddress(std::string(target[i].srcIp));
        input[i].dip = HcclIpAddress(std::string(target[i].dstIp));
        input[i].srcPort = target[i].srcPort;
        input[i].sl = target[i].sl;
        input[i].tc = target[i].tc;
        input[i].port = target[i].port;
        input[i].len = target[i].payloadLen;
        s32 sRet = memcpy_s(input[i].payload, input[i].len, target[i].payload, target[i].payloadLen);
        if (sRet != EOK) {
            HCCL_WARNING("[HccnRpingAddTarget]memcpy payload fail. errorno[%d] params:dstMaxSize[%u] dstPtr[%p] "
                         "srclen[%d] srcPayload[%p]",
                sRet, input[i].len, input[i].payload, target[0].payloadLen, target[0].payload);
            continue;
        }
    }
    ret = rping->HccnRpingAddTarget(static_cast<u32>(devLogicId), targetNum, input);
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[HccnRpingAddTarget]Device[%d] add targets fail", devLogicId);
        delete[] input;
        return HCCN_E_FAIL;
    }
 
    delete[] input;
    return HCCN_SUCCESS;
}
 
HccnResult HccnRpingRemoveTarget(HccnRpingCtx rpingCtx, uint32_t targetNum, HccnRpingTargetInfo *target)
{
    // 校验入参
    CHK_PRT_RET(target == nullptr, HCCL_ERROR("[HccnRpingRemoveTarget]target is null."), HCCN_E_PARA);
    CHK_PRT_RET(targetNum > TARGET_NUM_MAX,
        HCCL_ERROR("[HccnRpingRemoveTarget]targetNum[%u] is more than %u!", targetNum, TARGET_NUM_MAX), HCCN_E_PARA);
    // 校验指针，并将其转换为pingmesh指针
    CHK_PRT_RET(rpingCtx == nullptr, HCCL_ERROR("[HccnRpingRemoveTarget]rpingCtx is null."), HCCN_E_PARA);
    PingMesh *rping = static_cast<PingMesh*>(rpingCtx);
    // 获取device id
    s32 devLogicId = 0;
    HcclResult ret = hrtGetDeviceRefresh(&devLogicId);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingRemoveTarget]cannot get device logic id."), HCCN_E_FAIL);

    // 判断devLogicId与pingmesh所记录的是否一致
    if (devLogicId != rping->GetDeviceLogicId()) {
        HCCL_ERROR("[HccnRpingRemoveTarget]curr devId[%d] don't match record logicId[%d].",
            devLogicId, rping->GetDeviceLogicId());
        return HCCN_E_PARA;
    }
 
    // 将入参转换为内部接口可以使用的类型
    RpingInput *input = new (std::nothrow) RpingInput[targetNum];
    CHK_PRT_RET(input == nullptr, HCCL_ERROR("[HccnRpingRemoveTarget]memory alloc failed."), HCCN_E_FAIL);
    for (uint32_t i = 0; i < targetNum; i++) {
        input[i].sip = HcclIpAddress(std::string(target[i].srcIp));
        input[i].dip = HcclIpAddress(std::string(target[i].dstIp));
        input[i].sl = target[i].sl;
        input[i].tc = target[i].tc;
    }
 
    ret = rping->HccnRpingRemoveTarget(static_cast<u32>(devLogicId), targetNum, input);
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[HccnRpingRemoveTarget]Device[%d] remove targets fail", devLogicId);
        delete[] input;
        return HCCN_E_FAIL;
    }
 
    delete[] input;
    return HCCN_SUCCESS;
}

inline void ConvertTargetState(uint32_t targetNum, int *state, HccnRpingAddTargetState *targetState)
{
    for (uint32_t i = 0; i < targetNum; i++) {
        switch(static_cast<RpingLinkState>(state[i])) {
            case RpingLinkState::CONNECTED :
                targetState[i] = HCCN_RPING_ADDTARGET_STATE_DONE;
                break;
            case RpingLinkState::CONNECTING :
                targetState[i] = HCCN_RPING_ADDTARGET_STATE_DOING;
                break;
            case RpingLinkState::DISCONNECTED :
                targetState[i] = HCCN_RPING_ADDTARGET_STATE_FAIL;
                break;
            case RpingLinkState::TIMEOUT :
                targetState[i] = HCCN_RPING_ADDTARGET_STATE_TIMEOUT;
                break;
            default:
                targetState[i] = HCCN_RPING_ADDTARGET_STATE_RESERVED;
                break;
        }
    }
}
 
HccnResult HccnRpingGetTarget(HccnRpingCtx rpingCtx, uint32_t targetNum, HccnRpingTargetInfo *target,
                              HccnRpingAddTargetState *targetState)
{
    // 校验入参
    CHK_PRT_RET(target == nullptr, HCCL_ERROR("[HccnRpingGetTarget]target is null."), HCCN_E_PARA);
    CHK_PRT_RET(targetState == nullptr, HCCL_ERROR("[HccnRpingGetTarget]targetState is null."), HCCN_E_PARA);
    CHK_PRT_RET(targetNum > TARGET_NUM_MAX,
        HCCL_ERROR("[HccnRpingGetTarget]targetNum[%u] is more than %u!", targetNum, TARGET_NUM_MAX), HCCN_E_PARA);
    // 校验指针，并将其转换为pingmesh指针
    CHK_PRT_RET(rpingCtx == nullptr, HCCL_ERROR("[HccnRpingGetTarget]rpingCtx is null."), HCCN_E_PARA);
    PingMesh *rping = static_cast<PingMesh*>(rpingCtx);
    // 获取device id
    s32 devLogicId = 0;
    HcclResult ret = hrtGetDeviceRefresh(&devLogicId);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingGetTarget]cannot get device logic id."), HCCN_E_FAIL);

    // 判断devLogicId与pingmesh所记录的是否一致
    if (devLogicId != rping->GetDeviceLogicId()) {
        HCCL_ERROR("[HccnRpingGetTarget]curr devId[%d] don't match record logicId[%d].",
            devLogicId, rping->GetDeviceLogicId());
        return HCCN_E_PARA;
    }
 
    // 将入参转换为内部接口可以使用的类型
    RpingInput *input = new (std::nothrow) RpingInput[targetNum];
    CHK_PRT_RET(input == nullptr, HCCL_ERROR("[HccnRpingGetTarget]memory alloc failed."), HCCN_E_MEM);
    for (uint32_t i = 0; i < targetNum; i++) {
        input[i].sip = HcclIpAddress(std::string(target[i].srcIp));
        input[i].dip = HcclIpAddress(std::string(target[i].dstIp));
    }

    int *state = new (std::nothrow) int[targetNum];
    CHK_PRT_RET(state == nullptr, HCCL_ERROR("[HccnRpingGetTarget]memory alloc failed."), HCCN_E_MEM);
    ret = rping->HccnRpingGetTarget(static_cast<u32>(devLogicId), targetNum, input, state);
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[HccnRpingGetTarget]Device[%d] get targets fail", devLogicId);
        delete[] input;
        delete[] state;
        return HCCN_E_FAIL;
    }

    ConvertTargetState(targetNum, state, targetState);
    delete[] input;
    delete[] state;
    return HCCN_SUCCESS;
}
 
HccnResult HccnRpingBatchPingStart(HccnRpingCtx rpingCtx, uint32_t pktNum, uint32_t interval, uint32_t timeout)
{
    // 校验指针，并将其转换为pingmesh指针
    CHK_PRT_RET(rpingCtx == nullptr, HCCL_ERROR("[HccnRpingBatchPingStart]rpingCtx is null.", rpingCtx), HCCN_E_PARA);
    PingMesh *rping = static_cast<PingMesh*>(rpingCtx);
    // 获取device id
    s32 devLogicId = 0;
    HcclResult ret = hrtGetDeviceRefresh(&devLogicId);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingBatchPingStart]cannot get device logic id."), HCCN_E_PARA);

    // 判断devLogicId与pingmesh所记录的是否一致
    if (devLogicId != rping->GetDeviceLogicId()) {
        HCCL_ERROR("[HccnRpingBatchPingStart]curr devId[%d] don't match record logicId[%d].",
            devLogicId, rping->GetDeviceLogicId());
        return HCCN_E_PARA;
    }

    CHK_PRT_RET(rping->HccnRpingBatchPingStart(static_cast<u32>(devLogicId), pktNum, interval, timeout) != HCCL_SUCCESS,
        HCCL_ERROR("[HccnRpingBatchPingStart]task start failed."), HCCN_E_FAIL);

    return HCCN_SUCCESS;
}
 
HccnResult HccnRpingBatchPingStop(HccnRpingCtx rpingCtx)
{
    // 校验指针，并将其转换为pingmesh指针
    CHK_PRT_RET(rpingCtx == nullptr, HCCL_ERROR("[HccnRpingBatchPingStop]rpingCtx is null.", rpingCtx), HCCN_E_PARA);
    PingMesh *rping = static_cast<PingMesh*>(rpingCtx);
    // 获取device id
    s32 devLogicId = 0;
    HcclResult ret = hrtGetDeviceRefresh(&devLogicId);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingBatchPingStop]cannot get device logic id."), HCCN_E_PARA);

    // 判断devLogicId与pingmesh所记录的是否一致
    if (devLogicId != rping->GetDeviceLogicId()) {
        HCCL_ERROR("[HccnRpingBatchPingStop]curr devId[%d] don't match record logicId[%d].",
            devLogicId, rping->GetDeviceLogicId());
        return HCCN_E_PARA;
    }

    CHK_PRT_RET(rping->HccnRpingBatchPingStop(static_cast<u32>(devLogicId)) != HCCL_SUCCESS,
        HCCL_ERROR("[HccnRpingBatchPingStop]task stop failed."), HCCN_E_FAIL);

    return HCCN_SUCCESS;
}

inline void ConvertResultState(uint32_t state, HccnRpingResultState &resultState)
{
    if (state == RPING_RESULT_STATE_VALID) {
        resultState = HCCN_RPING_RESULT_STATE_VALID;
    } else {
        resultState = HCCN_RPING_RESULT_STATE_INVALID;
    }
}
 
HccnResult HccnRpingGetResult(HccnRpingCtx rpingCtx, uint32_t targetNum, HccnRpingTargetInfo *target,
                              HccnRpingResultInfo *result)
{
    // 校验入参
    CHK_PRT_RET(target == nullptr, HCCL_ERROR("[HccnRpingGetResult]target is null."), HCCN_E_PARA);
    CHK_PRT_RET(result == nullptr, HCCL_ERROR("[HccnRpingGetResult]result is null."), HCCN_E_PARA);
    CHK_PRT_RET(targetNum > TARGET_NUM_MAX,
        HCCL_ERROR("[HccnRpingGetResult]targetNum[%u] is more than %u!", targetNum, TARGET_NUM_MAX), HCCN_E_PARA);
    // 校验指针，并将其转换为pingmesh指针
    CHK_PRT_RET(rpingCtx == nullptr, HCCL_ERROR("[HccnRpingGetResult]rpingCtx is null."), HCCN_E_PARA);
    PingMesh *rping = static_cast<PingMesh*>(rpingCtx);
    // 获取device id
    s32 devLogicId = 0;
    HcclResult ret = hrtGetDeviceRefresh(&devLogicId);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingGetResult]cannot get device logic id."), HCCN_E_FAIL);

    // 判断devLogicId与pingmesh所记录的是否一致
    if (devLogicId != rping->GetDeviceLogicId()) {
        HCCL_ERROR("[HccnRpingGetResult]curr devId[%d] not match record[%d].", devLogicId, rping->GetDeviceLogicId());
        return HCCN_E_PARA;
    }
 
    // 将入参转换为内部接口可以使用的类型
    RpingInput *input = new (std::nothrow) RpingInput[targetNum];
    CHK_PRT_RET(input == nullptr, HCCL_ERROR("[HccnRpingGetResult]memory alloc failed."), HCCN_E_MEM);
    RpingOutput *output = new (std::nothrow) RpingOutput[targetNum];
    CHK_PRT_RET(output == nullptr, HCCL_ERROR("[HccnRpingGetResult]memory alloc failed."), HCCN_E_MEM);
    for (uint32_t i = 0; i < targetNum; i++) {
        input[i].sip = HcclIpAddress(std::string(target[i].srcIp));
        input[i].dip = HcclIpAddress(std::string(target[i].dstIp));
    }

    ret = rping->HccnRpingGetResult(static_cast<u32>(devLogicId), targetNum, input, output);
    if (ret == HCCL_E_AGAIN) {
        delete[] input;
        delete[] output;
        return HCCN_E_AGAIN;
    }
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[HccnRpingGetResult]Device[%d] get result fail", devLogicId);
        delete[] input;
        delete[] output;
        return HCCN_E_FAIL;
    }
 
    for (uint32_t i = 0; i < targetNum; i++) {
        ConvertResultState(output[i].state, result[i].state);
        if (output[i].state != ping_result_state::PING_RESULT_STATE_VALID) {
            HCCL_WARNING("[HccnRpingGetResult]Target[%d]'s state is not valid.", i);
            continue;
        }
        result[i].txPkt = output[i].txPkt;
        result[i].rxPkt = output[i].rxPkt;
        result[i].minRTT = output[i].minRTT;
        result[i].maxRTT = output[i].maxRTT;
        result[i].avgRTT = output[i].avgRTT;
    }
 
    delete[] input;
    delete[] output;
    return HCCN_SUCCESS;
}

HccnResult HccnRpingGetPayload(HccnRpingCtx rpingCtx, void **payload, uint32_t *payloadLen)
{
    // 校验入参
    CHK_PRT_RET(payload == nullptr, HCCL_ERROR("[HccnRpingGetPayload]payload is null."), HCCN_E_PARA);
    CHK_PRT_RET(payloadLen == nullptr, HCCL_ERROR("[HccnRpingGetPayload]payloadLen is null."), HCCN_E_PARA);
    // 校验指针，并将其转换为pingmesh指针
    CHK_PRT_RET(rpingCtx == nullptr, HCCL_ERROR("[HccnRpingGetPayload]rpingCtx is null."), HCCN_E_PARA);
    PingMesh *rping = static_cast<PingMesh*>(rpingCtx);
    // 获取device id
    s32 devLogicId = 0;
    HcclResult ret = hrtGetDeviceRefresh(&devLogicId);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingGetPayload]cannot get device logic id."), HCCN_E_FAIL);

    // 判断devLogicId与pingmesh所记录的是否一致
    if (devLogicId != rping->GetDeviceLogicId()) {
        HCCL_ERROR("[HccnRpingGetResult]curr devId[%d] not match record[%d].", devLogicId, rping->GetDeviceLogicId());
        return HCCN_E_PARA;
    }

    ret = rping->HccnRpingGetPayload(static_cast<u32>(devLogicId), payload, payloadLen);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[HccnRpingGetPayload]Device[%d] get payload fail", devLogicId),
        HCCN_E_FAIL);

    return HCCN_SUCCESS;
}

#ifdef __cplusplus
}
#endif // __cplusplus